import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules

# 加载数据
df = pd.read_excel('meal_order_detail.xlsx')

# 根据下单时间，转换为午餐/晚餐信息
# 午餐为0，晚餐为1
def get_meal_type(order_time):
    hour = order_time.hour
    return 0 if hour < 16 else 1

df['place_order_time'] = pd.to_datetime(df['place_order_time'])
df['meal_type'] = df['place_order_time'].apply(get_meal_type)

# 创建事务数据
transactions = df.groupby('order_id').apply(lambda x: x['dishes_id'].tolist() + [x['meal_type'].iloc[0]])

transactions.to_csv('transactions.csv', index=False)
transactions = transactions.to_list()

# 使用TransactionEncoder转换数据
te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)
df_trans = pd.DataFrame(te_ary, columns=te.columns_)

# 应用Apriori算法找出频繁项集
frequent_itemsets = apriori(df_trans, min_support=0.01, use_colnames=True)

# 生成关联规则
rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=0.5)

# 输出关联规则
rules.to_csv('association_rules.csv', index=False)
